import torch
import numpy as np
from typing import Tuple, List
from torch.utils.data import Dataset, DataLoader

class TimeSeriesDataset(Dataset):
    def __init__(self, data: np.ndarray, sequence_length: int, target_horizon: int = 1):
        self.data = torch.FloatTensor(data)
        self.sequence_length = sequence_length
        self.target_horizon = target_horizon
        
    def __len__(self):
        return len(self.data) - self.sequence_length - self.target_horizon + 1
    
    def __getitem__(self, idx):
        x = self.data[idx:idx + self.sequence_length]
        y = self.data[idx + self.sequence_length:idx + self.sequence_length + self.target_horizon]
        return x, y

def prepare_data_loaders(
    data: np.ndarray,
    batch_size: int,
    sequence_length: int,
    train_ratio: float = 0.8
) -> Tuple[DataLoader, DataLoader]:
    """准备训练集和验证集的数据加载器"""
    # 划分训练集和验证集
    train_size = int(len(data) * train_ratio)
    train_data = data[:train_size]
    val_data = data[train_size:]
    
    # 创建数据集
    train_dataset = TimeSeriesDataset(train_data, sequence_length)
    val_dataset = TimeSeriesDataset(val_data, sequence_length)
    
    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    
    return train_loader, val_loader

def save_model(model: torch.nn.Module, path: str):
    """保存模型"""
    torch.save(model.state_dict(), path)

def load_model(model: torch.nn.Module, path: str):
    """加载模型"""
    model.load_state_dict(torch.load(path))
    return model 